import os
import subprocess
from pathlib import Path
import re
import numpy as np
from pdbench.utils.data_utils import extract_tm_score


if __name__ == "__main__":
    native_pdb_path = "/home/bingwu/Downloads/Conformation_Generation_Project/Dynamics/T1024_GT/native/T1024.pdb"
    target_path = "/home/bingwu/Downloads/Conformation_Generation_Project/Dynamics/T1024_bioemu_pdb_sequence_backups/"
    target_mapping = {id: os.path.join(target_path, id) for id in os.listdir(target_path) }
    tm_exec_path = "/home/bingwu/Softwares/TMscore"

    total_tm_scores = []
    for f in target_mapping:
        cmd = [tm_exec_path, native_pdb_path, str(target_mapping[f])]
        print(str(target_mapping[f]))
        try:
            output = subprocess.check_output(cmd, text=True)
            tm_score = extract_tm_score(output)
            total_tm_scores.append(tm_score)
            print(f"{f}: {tm_score}")
        except subprocess.CalledProcessError as e:
            print(f"Failed to run TMscore for {f}: {e}")

    tm_score_data = np.array(total_tm_scores)
    print(tm_score_data)
    print(np.mean(tm_score_data), np.max(tm_score_data), np.min(tm_score_data), np.median(tm_score_data))

    # T1024_new
    # [0.3814 0.3794 0.3831 ... 0.3877 0.3804 0.3833]
    # 0.3825263 0.4228 0.3165 0.3826

    # [0.3771 0.3831 0.3924 ... 0.3786 0.3835 0.3851]
    # 0.3825263 0.4228 0.3165 0.3826

    # [0.6811 0.6874 0.6399 ... 0.4782 0.6558 0.6946]
    # 0.6305790333333333 0.753 0.3857 0.64455